#!/bin/bash
"""
DeepSpeed MOE + Qwen3-Coder API系统部署脚本
用于自动化部署和配置整个系统
"""

set -e  # 遇到错误时退出

# 颜色输出
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

# 日志函数
log_info() {
    echo -e "${GREEN}[INFO]${NC} $1"
}

log_warn() {
    echo -e "${YELLOW}[WARN]${NC} $1"
}

log_error() {
    echo -e "${RED}[ERROR]${NC} $1"
}

log_step() {
    echo -e "${BLUE}[STEP]${NC} $1"
}

# 检查Python版本
check_python() {
    log_step "检查Python环境..."
    
    if ! command -v python3 &> /dev/null; then
        log_error "Python3 未安装，请先安装Python 3.8+"
        exit 1
    fi
    
    python_version=$(python3 --version | cut -d' ' -f2 | cut -d'.' -f1-2)
    required_version="3.8"
    
    if [ "$(printf '%s\n' "$required_version" "$python_version" | sort -V | head -n1)" != "$required_version" ]; then
        log_error "Python版本过低，需要Python 3.8+，当前版本: $python_version"
        exit 1
    fi
    
    log_info "Python版本检查通过: $python_version"
}

# 检查CUDA
check_cuda() {
    log_step "检查CUDA环境..."
    
    if ! command -v nvidia-smi &> /dev/null; then
        log_warn "NVIDIA驱动未安装或CUDA不可用，将使用CPU模式"
        USE_CPU=true
        return
    fi
    
    cuda_version=$(nvidia-smi --query-gpu=driver_version --format=csv,noheader,nounits | head -1)
    log_info "CUDA驱动版本: $cuda_version"
    
    # 检查GPU内存
    gpu_memory=$(nvidia-smi --query-gpu=memory.total --format=csv,noheader,nounits | head -1)
    if [ "$gpu_memory" -lt 20000 ]; then
        log_warn "GPU内存不足 (${gpu_memory}MB)，建议至少20GB内存"
        log_warn "系统将使用CPU模式"
        USE_CPU=true
    else
        log_info "GPU内存充足: ${gpu_memory}MB"
        USE_CPU=false
    fi
}

# 安装系统依赖
install_system_deps() {
    log_step "安装系统依赖..."
    
    if command -v apt-get &> /dev/null; then
        sudo apt-get update
        sudo apt-get install -y build-essential cmake git curl wget
    elif command -v yum &> /dev/null; then
        sudo yum install -y gcc gcc-c++ make cmake git curl wget
    elif command -v dnf &> /dev/null; then
        sudo dnf install -y gcc gcc-c++ make cmake git curl wget
    else
        log_warn "无法识别包管理器，请手动安装系统依赖"
    fi
    
    log_info "系统依赖安装完成"
}

# 创建虚拟环境
setup_venv() {
    log_step "设置Python虚拟环境..."
    
    if [ ! -d "venv" ]; then
        python3 -m venv venv
        log_info "虚拟环境创建完成"
    else
        log_info "虚拟环境已存在"
    fi
    
    source venv/bin/activate
    pip install --upgrade pip setuptools wheel
    
    log_info "虚拟环境已激活"
}

# 安装Python依赖
install_python_deps() {
    log_step "安装Python依赖包..."
    
    source venv/bin/activate
    
    # 安装PyTorch（CUDA版本优先）
    if [ "$USE_CPU" = false ]; then
        pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118
    else
        pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu
    fi
    
    # 安装DeepSpeed
    pip install deepspeed
    
    # 安装其他依赖
    pip install fastapi uvicorn transformers accelerate datasets
    
    # 安装API服务相关依赖
    pip install python-multipart aiofiles python-jose[cryptography] passlib[bcrypt]
    
    # 安装监控和日志依赖
    pip install prometheus-client psutil
    
    log_info "Python依赖安装完成"
}

# 下载模型（可选）
download_model() {
    log_step "检查模型文件..."
    
    MODEL_PATH=${MODEL_PATH:-"/workspace/models/qwen3-coder-480b-a35b-instruct"}
    
    if [ ! -d "$MODEL_PATH" ]; then
        log_info "模型目录不存在: $MODEL_PATH"
        log_info "请手动下载Qwen3-Coder模型到指定目录"
        log_info "模型下载命令: "
        log_info "  git clone https://huggingface.co/Qwen/Qwen3-Coder-480B-A35B-Instruct $MODEL_PATH"
    else
        log_info "模型目录存在: $MODEL_PATH"
    fi
}

# 生成配置文件
generate_configs() {
    log_step "生成配置文件..."
    
    source venv/bin/activate
    
    # 生成Roo Code配置
    python roo_code_config.py
    
    # 创建环境变量文件
    cat > .env << EOF
# DeepSpeed MOE + Qwen3-Coder API配置
API_HOST=0.0.0.0
API_PORT=8000
ENVIRONMENT=development
DEBUG=true

# 模型配置
MODEL_PATH=$MODEL_PATH
TENSOR_PARALLEL_SIZE=8
EXPERT_PARALLEL_SIZE=8
NUM_EXPERTS=128

# GPU配置
USE_CPU=$USE_CPU

# API Keys
ROO_CODE_KEY=roo-code-key-2024
CLAUDE_CODE_KEY=claude-code-key-2024
LOCAL_DEV_KEY=local-dev-key-2024
EOF
    
    log_info "配置文件生成完成"
}

# 创建启动脚本
create_startup_scripts() {
    log_step "创建启动脚本..."
    
    # 创建启动API服务器的脚本
    cat > start_api.sh << 'EOF'
#!/bin/bash
source venv/bin/activate

# 检查环境变量
if [ ! -z "$MODEL_PATH" ]; then
    export MODEL_PATH=$MODEL_PATH
fi

# 启动API服务器
echo "启动DeepSpeed MOE + Qwen3-Coder API服务器..."
echo "访问地址: http://localhost:8000"
echo "API文档: http://localhost:8000/docs"
echo "健康检查: http://localhost:8000/health"

python api_server.py
EOF
    
    chmod +x start_api.sh
    
    # 创建停止脚本
    cat > stop_api.sh << 'EOF'
#!/bin/bash
echo "停止API服务器..."
pkill -f "api_server.py"
pkill -f "uvicorn"
echo "API服务器已停止"
EOF
    
    chmod +x stop_api.sh
    
    # 创建测试脚本
    cat > test_api.sh << 'EOF'
#!/bin/bash
echo "测试API服务器..."

# 测试健康检查
echo "1. 健康检查:"
curl -s http://localhost:8000/health | jq .

echo -e "\n2. 获取模型列表:"
curl -s -H "Authorization: Bearer roo-code-key-2024" \
     http://localhost:8000/v1/models | jq .

echo -e "\n3. 测试聊天完成:"
curl -s -X POST http://localhost:8000/v1/chat/completions \
  -H "Content-Type: application/json" \
  -H "Authorization: Bearer roo-code-key-2024" \
  -d '{
    "model": "qwen3-coder-480b-a35b-instruct",
    "messages": [
      {"role": "user", "content": "Write a simple Python function to calculate factorial"}
    ],
    "max_tokens": 100
  }' | jq .

echo -e "\nAPI测试完成"
EOF
    
    chmod +x test_api.sh
    
    log_info "启动脚本创建完成"
}

# 创建Docker文件（可选）
create_docker_files() {
    log_step "创建Docker文件..."
    
    # Dockerfile
    cat > Dockerfile << 'EOF'
FROM nvidia/cuda:11.8-devel-ubuntu20.04

# 设置环境变量
ENV DEBIAN_FRONTEND=noninteractive
ENV PYTHONPATH=/app
ENV PYTHONUNBUFFERED=1

# 安装系统依赖
RUN apt-get update && apt-get install -y \
    python3.8 python3.8-pip python3.8-dev \
    build-essential cmake git curl wget \
    && rm -rf /var/lib/apt/lists/*

# 创建符号链接
RUN ln -sf /usr/bin/python3.8 /usr/bin/python

# 安装Python依赖
COPY requirements.txt /app/
RUN pip install --no-cache-dir -r /app/requirements.txt

# 复制应用代码
COPY . /app/

# 创建非root用户
RUN useradd -m -u 1000 appuser && chown -R appuser:appuser /app
USER appuser

# 暴露端口
EXPOSE 8000

# 启动命令
CMD ["python", "api_server.py"]
EOF
    
    # requirements.txt
    cat > requirements.txt << 'EOF'
torch>=2.0.0
deepspeed>=0.11.0
transformers>=4.30.0
fastapi>=0.100.0
uvicorn>=0.20.0
accelerate>=0.20.0
datasets>=2.10.0
python-multipart>=0.0.6
aiofiles>=23.0.0
python-jose[cryptography]>=3.3.0
passlib[bcrypt]>=1.7.4
prometheus-client>=0.17.0
psutil>=5.9.0
EOF
    
    # docker-compose.yml
    cat > docker-compose.yml << 'EOF'
version: '3.8'

services:
  deepspeed-qwen3-api:
    build: .
    ports:
      - "8000:8000"
    environment:
      - MODEL_PATH=/workspace/models/qwen3-coder-480b-a35b-instruct
      - USE_CPU=false
    volumes:
      - ./models:/workspace/models
      - ./logs:/app/logs
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              count: 1
              capabilities: [gpu]
    
  redis:
    image: redis:alpine
    ports:
      - "6379:6379"
    volumes:
      - redis_data:/data

volumes:
  redis_data:
EOF
    
    log_info "Docker文件创建完成"
}

# 创建systemd服务文件
create_systemd_service() {
    log_step "创建systemd服务文件..."
    
    sudo tee /etc/systemd/system/deepspeed-qwen3-api.service > /dev/null << EOF
[Unit]
Description=DeepSpeed MOE + Qwen3-Coder API Server
After=network.target

[Service]
Type=simple
User=$USER
WorkingDirectory=$(pwd)
Environment=PATH=$(pwd)/venv/bin
ExecStart=$(pwd)/venv/bin/python api_server.py
Restart=always
RestartSec=10
StandardOutput=journal
StandardError=journal

[Install]
WantedBy=multi-user.target
EOF
    
    sudo systemctl daemon-reload
    log_info "systemd服务文件已创建"
    log_info "启用服务: sudo systemctl enable deepspeed-qwen3-api"
    log_info "启动服务: sudo systemctl start deepspeed-qwen3-api"
    log_info "查看状态: sudo systemctl status deepspeed-qwen3-api"
}

# 主部署流程
main() {
    echo "==============================================="
    echo "   DeepSpeed MOE + Qwen3-Coder API部署脚本"
    echo "==============================================="
    
    # 解析命令行参数
    while [[ $# -gt 0 ]]; do
        case $1 in
            --cpu)
                USE_CPU=true
                shift
                ;;
            --model-path)
                MODEL_PATH="$2"
                shift 2
                ;;
            --help|-h)
                echo "用法: $0 [选项]"
                echo "选项:"
                echo "  --cpu              使用CPU模式"
                echo "  --model-path PATH  指定模型路径"
                echo "  --systemd          创建systemd服务"
                echo "  --docker           创建Docker文件"
                echo "  --help             显示帮助信息"
                exit 0
                ;;
            --systemd)
                CREATE_SYSTEMD=true
                shift
                ;;
            --docker)
                CREATE_DOCKER=true
                shift
                ;;
            *)
                log_error "未知参数: $1"
                exit 1
                ;;
        esac
    done
    
    # 执行部署步骤
    check_python
    check_cuda
    install_system_deps
    setup_venv
    install_python_deps
    download_model
    generate_configs
    create_startup_scripts
    
    if [ "$CREATE_DOCKER" = true ]; then
        create_docker_files
    fi
    
    if [ "$CREATE_SYSTEMD" = true ]; then
        create_systemd_service
    fi
    
    echo
    echo "==============================================="
    echo "              部署完成！"
    echo "==============================================="
    echo
    echo "接下来的步骤："
    echo "1. 下载Qwen3-Coder模型:"
    echo "   git clone https://huggingface.co/Qwen/Qwen3-Coder-480B-A35B-Instruct $MODEL_PATH"
    echo
    echo "2. 激活虚拟环境:"
    echo "   source venv/bin/activate"
    echo
    echo "3. 启动API服务器:"
    echo "   ./start_api.sh"
    echo
    echo "4. 测试API:"
    echo "   ./test_api.sh"
    echo
    echo "5. 配置Roo Code:"
    echo "   - 打开Roo Code设置"
    echo "   - 选择'OpenAI兼容'提供商"
    echo "   - 设置Base URL: http://localhost:8000"
    echo "   - 设置API Key: roo-code-key-2024"
    echo "   - 选择模型: qwen3-coder-480b-a35b-instruct"
    echo
    echo "访问地址:"
    echo "  API服务器: http://localhost:8000"
    echo "  API文档:   http://localhost:8000/docs"
    echo "  健康检查:  http://localhost:8000/health"
    echo
}

# 执行主函数
main "$@"